package com.ls.tree.huffmancode;

import java.io.*;
import java.util.*;

/**
 * 霍夫曼编码
 * 用于压缩的
 */
public class HuffmanCode {
    public static void main(String[] args) {
        HuffmanCode huffmanCode = new HuffmanCode();
        String s = "i like like like java do you like a java wd";
        byte[] bytes = s.getBytes();
//        System.out.println(bytes.length);
//        List<Node> list = huffmanCode.getNode(bytes);
//        Node huffmanTree = huffmanCode.createHuffmanTree(list);
//        huffmanTree.preOrder();
//        Map<Byte, String> codes = huffmanCode.getCodes(huffmanTree);
//        System.out.println(codes);
//        byte[] zip = huffmanCode.zip(bytes, codes);
//        System.out.println(zip.length);

//        byte[] huffmanZip = huffmanCode.huffmanZip(bytes);
//        System.out.println(Arrays.toString(huffmanZip));
//        System.out.println(huffmanCodes);
//        byte[] decode = huffmanCode.decode(huffmanCodes, huffmanZip);
//        System.out.println(new String(decode));

        huffmanCode.zipFile("E:\\s\\00.jpg", "E:\\t\\03.z");
        huffmanCode.unZipFile("E:\\t\\03.z", "E:\\t\\003.jpg");

    }

    /**
     * @param zipFile 源文件全路径
     * @param dstFile 目标路径
     */
    public void unZipFile(String zipFile, String dstFile) {
        try (FileInputStream is = new FileInputStream(zipFile);
             FileOutputStream os = new FileOutputStream(dstFile);
             ObjectInputStream ois = new ObjectInputStream(is)) {

            byte[] huffmanBytes = (byte[]) ois.readObject();
            Map<Byte, String> huffmanCodes = (Map<Byte, String>) ois.readObject();

            byte[] bytes = decode(huffmanCodes, huffmanBytes);

            os.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * @param srcFile 源文件全路径
     * @param dstFile 目标路径
     */
    public void zipFile(String srcFile, String dstFile) {
        try (FileInputStream is = new FileInputStream(srcFile);
             FileOutputStream os = new FileOutputStream(dstFile);
             ObjectOutputStream oos = new ObjectOutputStream(os)) {

            byte[] b = new byte[is.available()];
            is.read(b);

            //直接对源文件压缩
            byte[] huffmanBytes = huffmanZip(b);
            //这里我们以对象流的方式写入赫夫曼码,是为了以后恢复时使用
            oos.writeObject(huffmanBytes);
            oos.writeObject(huffmanCodes);


        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    //-----------------------
    public byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            byte b = huffmanBytes[i];
            boolean flag = (i == huffmanBytes.length - 1);
            stringBuilder.append(byteToBitString(!flag, b));
        }
        System.out.println(stringBuilder.toString());
        //调转键值
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }
        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length(); ) {
            int count = 1;
            boolean flag = true;
            Byte b = null;
            while (flag) {
                String key = stringBuilder.substring(i, i + count);
                b = map.get(key);
                if (b == null) {
                    count++;
                } else {
                    flag = false;
                }
            }
            list.add(b);
            i += count;
        }

        byte[] b = new byte[list.size()];
        for (int i = 0; i < b.length; i++) {
            b[i] = list.get(i);
        }
        return b;


    }

    /**
     * @param flag 标志是否需要补高位,如果是true表示需要补高位,如果是false表示不补,如果是最后一位,无需补高位
     * @param b
     * @return
     */
    public String byteToBitString(boolean flag, byte b) {
        int temp = b;
        if (flag) {//如果是正数我们还需补高位
            temp |= 256;
        }
        String s = Integer.toBinaryString(temp);//返回的是temp对应的二进制补码
        if (flag) {
            return s.substring(s.length() - 8);
        }
        return s;
    }


    public byte[] huffmanZip(byte[] bytes) {
        List<Node> list = getNode(bytes);
        Node huffmanTree = createHuffmanTree(list);
        Map<Byte, String> codes = getCodes(huffmanTree);
        byte[] zip = zip(bytes, codes);
        return zip;
    }

    /**
     * 用霍夫曼编码压缩
     *
     * @param bytes        原始byte[]
     * @param huffmanCodes 码表
     * @return 返回霍夫曼编码处理后的byte[]
     */
    public byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(huffmanCodes.get(b));
        }
        System.out.println(sb.toString());
        //统计长度
        int len;
        if (sb.length() % 8 == 0) {
            len = sb.length() / 8;
        } else {
            len = sb.length() / 8 + 1;
        }

        byte[] huffman = new byte[len];
        int index = 0;
        for (int i = 0; i < sb.length(); i += 8) {
            String strByt;
            if (i + 8 > sb.length()) {
                strByt = sb.substring(i);
            } else {
                strByt = sb.substring(i, i + 8);
            }

            //这里有问题{32=01, 97=100, 100=11001, 101=1101, 105=101, 106=0000, 107=1110, 108=1111, 111=0001, 117=00110, 118=0010, 119=00111, 121=11000}

            huffman[index] = (byte) Integer.parseInt(strByt, 2);
            index++;
        }
        return huffman;

    }


    static Map<Byte, String> huffmanCodes = new HashMap<>();
    static StringBuilder stringBuilder = new StringBuilder();

    public Map<Byte, String> getCodes(Node root) {
        getCodes(root.left, "0", stringBuilder);
        getCodes(root.right, "1", stringBuilder);
        return huffmanCodes;
    }

    /**
     * //根据树生成码表
     *
     * @param node          传入节点
     * @param code          路径:左子节点0 右子节点1
     * @param stringBuilder
     */
    public void getCodes(Node node, String code, StringBuilder stringBuilder) {
        StringBuilder stringBuilder2 = new StringBuilder(stringBuilder);
        stringBuilder2.append(code);
        if (node != null) {
            if (node.data == null) {//非叶子节点
                getCodes(node.left, "0", stringBuilder2);
                getCodes(node.right, "1", stringBuilder2);
            } else {
                //叶子
                huffmanCodes.put(node.data, stringBuilder2.toString());
            }
        }
    }

    /**
     * 根据字节创建节点
     *
     * @param bytes
     * @return
     */
    public List<Node> getNode(byte[] bytes) {
        ArrayList<Node> nodes = new ArrayList<>();
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if (count == null) {
                counts.put(b, 1);
            } else {
                counts.put(b, count + 1);
            }
        }
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }

        return nodes;
    }

    /**
     * 根据带权重的节点创建霍夫曼树
     *
     * @param nodes
     * @return
     */
    public Node createHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            Collections.sort(nodes);
            Node left = nodes.get(0);
            Node right = nodes.get(1);
            Node parent = new Node(null, left.weight + right.weight);
            parent.left = left;
            parent.right = right;
            nodes.remove(left);
            nodes.remove(right);

            nodes.add(parent);
        }
        return nodes.get(0);
    }


}

class Node implements Comparable<Node> {
    Byte data;//存放数据本身,比如'a' -> 97
    int weight;
    Node left;
    Node right;


    public Node(Byte data, int weight) {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public int compareTo(Node o) {

        return this.weight - o.weight;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    //前序遍历
    public void preOrder() {
        System.out.println(this);
        if (this.left != null) {
            this.left.preOrder();
        }
        if (this.right != null) {
            this.right.preOrder();
        }
    }
}
